#!/usr/bin/perl
use strict;
use warnings;
use Fcntl 'SEEK_SET';

## get options

use Getopt::Long;
my %param = (
    source           => undef,
    target           => undef,
    align_directory  => undef,
    method           => undef, ## what method to use to select the labels at each iteration.
    step             => undef, ## begin with zero labels: each dataset includes "step" more labels.
    trainreserve     => undef, ## the number of labels that are to be used to test unsupervised learning
    testreserve      => undef, ## the number of labels to exclusively use for testing
    max_total        => undef, ## maximum total size of dataset
    prefix           => undef, ## prefix of bitext files (output: prefix.totalsize.labeledsize)
    );
GetOptions(
    "source=s"           => \$param{source},
    "target=s"           => \$param{target},
    "align_directory=s"  => \$param{align_directory},
    "method=s"           => \$param{method},
    "step=i"             => \$param{step},
    "trainreserve=i"      => \$param{trainreserve},
    "testreserve=i"      => \$param{testreserve},
    "max_total=i"        => \$param{max_total},
    "prefix=s"           => \$param{prefix},
    );
while(my ($param, $value) = each(%param)){
    die "$param undefined!" if !defined($value);
}

## load hand alignments

my %all_hand_alignments = ();

my $data_dir = "$param{align_directory}/data";
-d $data_dir or die "could not find data directory!";
my $anno_dir = "$param{align_directory}/annotations";
-d $anno_dir or die "could not find annotation directory!";

opendir(ANNO, $anno_dir) or die "could not open annotations directory!";
while(defined(my $entry = readdir(ANNO))){
    next if ($entry eq ".") || ($entry eq "..");
    my $anno_subdir = "$anno_dir/$entry";
    next if !(-d $anno_subdir);

    ## get corresponding pair..
    my $corresp = "$data_dir/$entry";
    -f $corresp or die "could not find corresponding data ($entry)!";
    open(F, "<$corresp") or die "could not open data!";
    defined(my $source = <F>) or die "..";
    chomp $source;
    defined(my $target = <F>) or die "..";
    chomp $target;
    close(F);

    ## for the purposes of bookkeeping ..
    my $source_normspace = &normalize($source);
    my $target_normspace = &normalize($target);

    ## for each annotator ..
    opendir(SUBANNO, $anno_subdir) or die "could not open subannotations directory!";
    while(defined(my $sub_entry = readdir(SUBANNO))){
	my $anno_entry = "$anno_subdir/$sub_entry";
	next if !(-f $anno_entry);
	
	## read in alignment ..
	my @annotation = ();
	open(F, "<$anno_entry") or die "could not open annotation!";
	defined(my $line = <F>) or last; ##die "um .. $anno_entry";
	close(F);
	foreach my $pair (split(/;/, $line)){
	    my ($source_j, $target_i) = split(/,/, $pair);

	    push(@annotation, [$source_j, $target_i]);
	}
	
	## store alignment .. (NOT BEING DONE PROPERLY!)
	my $key = "$source_normspace##$target_normspace";
	$all_hand_alignments{$key} = \@annotation;
	
	last; ## !!! for now we only care about the one annotator..
    }
    closedir(SUBANNO);
}
closedir(ANNO);

print "no hand aligns: " . scalar(keys(%all_hand_alignments)) . "\n";

## create test file with some number of randomly chosen alignments
srand(837485); ## for the purposes of the experiment we want the same thing each time
my $testset_filename = "$param{prefix}.testset";
open(TESTSET, ">$testset_filename") or die "um.";
my $no_test_labels = 0;
while($no_test_labels < $param{testreserve}){
    
    ## choose a random alignment, and remove from set of available labels ..
    my @keys = keys(%all_hand_alignments);
    my $key = $keys[int(rand(scalar(@keys)))];
    my $annotation_ref = $all_hand_alignments{$key};
    delete $all_hand_alignments{$key};

    ## print this label out to test set ..
    my @srctgt = split(/##/, $key);
    my $source = $srctgt[0];
    my $target = $srctgt[1];
    print TESTSET "1\n";   ## strictly speaking this is wrong
    print TESTSET "$source\n";
    print TESTSET "$target\n";
    print TESTSET "test ";
    foreach my $st_ref (@$annotation_ref){
	my $st_pair = "$st_ref->[0] $st_ref->[1]";
	print TESTSET "$st_pair ";
    }
    print TESTSET "\n";

    ## increment count of test labels ..
    $no_test_labels++;
}
close(TESTSET);

## also, "randomly" choose some alignments for testing unsupervised alignment
my %unchosen_hand_alignments = %all_hand_alignments;
my %aligns_for_unsupervised = ();
my $no_unsupervised_labels = 0;
my $giza_golden_filename = "$param{prefix}.$param{max_total}.$param{method}.giza.golden";
open(GIZA_GOLDEN, ">$giza_golden_filename");
while($no_unsupervised_labels < $param{trainreserve}){
    
    ## choose a random alignment ..
    my @keys = keys(%all_hand_alignments);
    my $key = $keys[int(rand(scalar(@keys)))];
    my $annotation_ref = $all_hand_alignments{$key};
    delete $unchosen_hand_alignments{$key};  ## don't pick this again.

    ## add to set of unsupervised training labels ..
    $aligns_for_unsupervised{$key} = $annotation_ref;

    ## also save for giza
    my @srctgt = split(/##/, $key);
    my $source = $srctgt[0];
    my $target = $srctgt[1];
    print GIZA_GOLDEN "$source\n";
    print GIZA_GOLDEN "$target\n";
    foreach my $st_ref (@$annotation_ref){
	my $st_pair = "$st_ref->[0] $st_ref->[1]";
	print GIZA_GOLDEN "$st_pair ";
    }
    print GIZA_GOLDEN "\n";
    
    ## increment count!
    $no_unsupervised_labels++;
}
close(GIZA_GOLDEN);

print "no hand aligns: " . scalar(keys(%all_hand_alignments)) . "\n";
print "unchosen hand aligns: " . scalar(keys(%unchosen_hand_alignments)) . "\n";

## add "step" more labels each iteration. begin with zero labels.
my %hand_alignments = ();
my $target_number_labels = 0;
my $at_beginning_alignments_available = scalar(keys(%unchosen_hand_alignments));
while($target_number_labels < $at_beginning_alignments_available){ #%all_hand_alignments))){

    ## reach quota of alignments that we want
    while(scalar(keys(%hand_alignments)) < $target_number_labels){
	if($param{method} eq "random"){  ## choose randomly from set that we haven't chosen yet ..
	    my @keys = keys(%unchosen_hand_alignments);
	    my $random_key = $keys[int(rand(scalar(@keys)))];
	    $hand_alignments{$random_key} = $unchosen_hand_alignments{$random_key};
	    delete $unchosen_hand_alignments{$random_key};
	}else{
	    die "unrecognized method..";
	}
    }

    ## open output file for this version of corpus!
    my $bitext_filename = "$param{prefix}.$param{max_total}.$param{method}.${target_number_labels}";
    open(OUTPUT, ">$bitext_filename") or die "um.";

    ## open corresponding GIZA++ input file
    my $giza_filename = "$param{prefix}.$param{max_total}.$param{method}.${target_number_labels}.giza";
    open(GIZA, ">$giza_filename") or die;
    my %giza_src_token_to_id = ();
    my %giza_src_id_to_token = ();
    my $c_src = 1;
    my %giza_src_id_count = ();
    my %giza_tgt_token_to_id = ();
    my %giza_tgt_id_to_token = ();
    my $c_tgt = 1;
    my %giza_tgt_id_count = ();

    ## print all alignments (with labels if they have been chosen)!
    foreach my $key (keys(%all_hand_alignments)){
	my @srctgt = split(/##/, $key);
	my $source = $srctgt[0];
	my $target = $srctgt[1];
	my $annotation_ref = $all_hand_alignments{$key};
	
	## data
	print OUTPUT "1\n";   ## strictly speaking this is wrong
	print OUTPUT "$source\n";
	print OUTPUT "$target\n";

	## AAAAND output to GIZA++ format file.
	print GIZA "1\n";  ##ditto
	print GIZA &gizaOutput($target, \%giza_tgt_token_to_id, \%giza_tgt_id_to_token, \%giza_tgt_id_count, \$c_tgt) . "\n"; ## reverse..
	print GIZA &gizaOutput($source, \%giza_src_token_to_id, \%giza_src_id_to_token, \%giza_src_id_count, \$c_src) . "\n";

	## something that we chose for testing unsupervised learning?
	if(defined($aligns_for_unsupervised{$key})){
	    print OUTPUT "test ";
	    foreach my $st_ref (@$annotation_ref){
		my $st_pair = "$st_ref->[0] $st_ref->[1]";
		print OUTPUT "$st_pair ";
	    }
	}else{
	
	    ## something that we chose for TRAINING unsupervised learning?
	    if(defined($hand_alignments{$key})){
		print OUTPUT "train ";
		foreach my $st_ref (@$annotation_ref){
		    my $st_pair = "$st_ref->[0] $st_ref->[1]";
		    print OUTPUT "$st_pair ";
		}
	    }else{
		print OUTPUT "%";
	    }
	}
	print OUTPUT "\n";
    }

    ## print bitext up to a certain size, but not items that we've already printed
    open(SOURCE, "<$param{source}") or die "could not open source file!";
    open(TARGET, "<$param{target}") or die "could not open target file!";
    my $printed = scalar(keys(%all_hand_alignments));
    while(($printed < $param{max_total}) && defined(my $source = <SOURCE>)){
	chomp $source;
	$source =~ s/^\s+(\S|$)/$1/;
	$source =~ s/(^|\S)\s+$/$1/;
	my @source_words = split(/\s+/, $source);
	
	## read corresponding target
	defined(my $target = <TARGET>) or die "lopsided bitext!";
	chomp $target;
	$target =~ s/^\s+(\S|$)/$1/;
	$target =~ s/(^|\S)\s+$/$1/;
	my @target_words = split(/\s+/, $target);

	## print if not already printed
	my $key = &normalize($source) . "##" . &normalize($target);
	if(!defined(my $annotation_ref = $all_hand_alignments{$key})){
	    print OUTPUT "1\n";   ## strictly speaking this is wrong
	    print OUTPUT "$source\n";
	    print OUTPUT "$target\n";
	    print OUTPUT "%\n";

	    print GIZA "1\n";  ##ditto
	    print GIZA &gizaOutput($target, \%giza_tgt_token_to_id, \%giza_tgt_id_to_token, \%giza_tgt_id_count, \$c_tgt) . "\n";  ##reverse
	    print GIZA &gizaOutput($source, \%giza_src_token_to_id, \%giza_src_id_to_token, \%giza_src_id_count, \$c_src) . "\n";
	    
	    $printed++;
	}
    }
    close(SOURCE);
    close(TARGET);

    ## this version of corpus is complete!!
    close(OUTPUT);
    close(GIZA);

    ## also dump GIZA++ lexicon..
    &dumpGIZAVocab(\%giza_src_id_to_token, \%giza_src_id_count,
		   "$param{prefix}.$param{max_total}.$param{method}.${target_number_labels}.giza.t.vocab");  ## reverse
    &dumpGIZAVocab(\%giza_tgt_id_to_token, \%giza_tgt_id_count,
		   "$param{prefix}.$param{max_total}.$param{method}.${target_number_labels}.giza.s.vocab");

    ## something for runscript
    my $cmd = "./wordalign++ -schedule test/schedule -bitext BITEXT -testset TESTSET";
    $cmd =~ s/BITEXT/${bitext_filename}/g;
    $cmd =~ s/TESTSET/${testset_filename}/g;
    print "$cmd\n";

    ## increase quota of labels that we want for next round
    $target_number_labels += $param{step};
}

sub dumpGIZAVocab{
    my ($inv_vocab_ref, $count_ref, $filename) = @_;

    open(F, ">$filename") or die;
    foreach my $id (sort {$a <=> $b} keys(%$inv_vocab_ref)){
	print F "$id\t$inv_vocab_ref->{$id}\t$count_ref->{$id}\n";
    }
    close(F);
}

sub gizaOutput{
    my ($s, $vocab_ref, $vocab_inv_ref, $count_ref, $c) = @_;
    my @ids = ();

    foreach my $token (split(/\s+/, $s)){
	if(!defined($vocab_ref->{$token})){
	    $vocab_ref->{$token} = $$c;
	    $vocab_inv_ref->{$$c} = $token;
	    $$c++;
	}
	$count_ref->{$vocab_ref->{$token}}++;
	push(@ids, $vocab_ref->{$token});
    }

    return join(" ", @ids);
}

sub normalize{
    my ($s) = @_;

    $s =~ s/(\s)\s+/$1/g;    

    return $s;
}
